import argparse
import os
import pandas as pd
import matplotlib.pyplot as plt

def plot_csv_files(directory):
    plt.figure()
    
    for filename in os.listdir(directory):
        if filename.endswith('.csv'):
            print(f'Plotting {filename}')
            filepath = os.path.join(directory, filename)
            data = pd.read_csv(filepath)
            label = os.path.splitext(filename)[0]
            plt.plot(data['MatrixSize'], data['GFLOPS'], label=label)
    
    plt.xlabel('Matrix Size')
    plt.ylabel('GFLOPS')
    plt.title('GFLOPS vs Matrix Size')
    plt.legend(loc='lower right')
    plt.savefig(os.path.join(directory, 'COMPARE_ALL.png'))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Plot GFLOPS vs Matrix Size from CSV files in a directory.")
    parser.add_argument('directory', type=str, help='Path to the directory containing CSV files')
    
    args = parser.parse_args()
    plot_csv_files(args.directory)



# Write a python script that iterates a directory indicated by a CLI argument, searching for all *.csv files, and plot one line based on each file, then label each line by the file name.
# Each csv file contains this content. The y-axis of the line is GFLOPS, while the x-axis is MatrixSize.
# ```
# MatrixSize, GFLOPS, Diff
# 1024, 8910.31, 8.010864e-05 
# 1280, 13607.97, 0.000000e+00 
# 1536, 13808.66, 1.220703e-04 
# 1792, 14454.69, 0.000000e+00 
# 2048, 15472.59, 0.000000e+00 
# 2304, 16463.40, 0.000000e+00 
# 2560, 15521.63, 0.000000e+00 
# 2816, 16798.92, 0.000000e+00 
# 3072, 16466.73, 0.000000e+00 
# 3328, 16452.50, 0.000000e+00 
# 3584, 16643.04, 0.000000e+00 
# 3840, 16959.77, 0.000000e+00 
# 4096, 17219.56, 0.000000e+00 
# ```


